from typing import Dict, List

from matplotlib import pyplot as plt

from util import *
from wall import Wall


def g(x: float):
    return 0 if x < 0 else x


class Simulation:
    A: float = 2000  # N
    B: float = 0.08  # m
    k: float = 120000  # kg/s
    K: float = 240000  # kg/ms

    time: int = 0
    time_gap: float = 0.005  # s
    agents_hashmap: Dict[int, Agent]
    walls: List[Wall]

    def __init__(self, agents_hashmap, walls):
        self.tao = 0.3
        self.agents_hashmap = agents_hashmap
        self.walls = walls

    def simulate(self):
        self.time += 1
        # self.change_agents_destination()
        self.cal_agents_force()
        self.move_agents()
        self.remove_agents()
        self.draw()

    def cal_agents_force(self):
        for _, agent in self.agents_hashmap.items():
            agent.force = self.cal_force(agent)

    def cal_force(self, agent: Agent):
        force = self.cal_self_driving_force(agent)
        force += self.cal_other_agents_force(agent)
        force += self.cal_walls_force(agent)
        return force

    def draw(self):
        self.draw_agents()
        self.draw_walls()

    def draw_agents(self):
        for _, agent in self.agents_hashmap.items():
            drawn_graphics = plt.Circle((agent.position[0], agent.position[1]),
                                        agent.radius - 0.03, color='y', fill='y')
            plt.gcf().gca().add_artist(drawn_graphics)

    def draw_walls(self):
        for wall in self.walls:
            plt.plot([wall.start_point[0], wall.end_point[0]],
                     [wall.start_point[1], wall.end_point[1]],
                     color='r', alpha=1)

    def cal_self_driving_force(self, agent: Agent):
        agent_expected_vel = np.array(agent.destination - agent.position)
        agent_expected_unit_vel = agent_expected_vel / np.linalg.norm(agent_expected_vel)
        self_driving_force = agent.m * (agent.expected_speed * agent_expected_unit_vel - agent.actual_speed) \
                             / self.tao
        # print("自驱力：", np.linalg.norm(self_driving_force))
        return self_driving_force

    def cal_other_agents_force(self, agent):
        other_agents_force = 0
        for name, other in self.agents_hashmap.items():
            if name == agent.name:
                continue
            other_agents_force += self.cal_other_force(agent, other)
        # print("个体力：", np.linalg.norm(other_agents_force))
        return other_agents_force

    def cal_walls_force(self, agent):
        walls_force = 0
        for wall in self.walls:
            walls_force += self.cal_wall_force(agent, wall)
        return walls_force

    def cal_other_force(self, agent: Agent, other: Agent):
        rij = agent.radius + other.radius
        dij = np.linalg.norm(agent.position - other.position)
        nij = (agent.position - other.position) / dij
        tij = np.array([-nij[1], nij[0]])
        v_diff = (other.actual_speed - agent.actual_speed) * tij
        x = rij - dij
        other_force = (self.A * np.exp(x / self.B) + self.k * g(x)) * nij \
                      + self.K * g(x) * v_diff * tij
        return other_force

    def cal_wall_force(self, agent: Agent, wall: Wall):
        v = get_distance_a_to_w(agent, wall)
        tiw = np.array([-v.niw[1], v.niw[0]])
        if not v.flg:
            return np.array([0.0, 0.0])
        x = agent.radius - v.diw
        wall_force = (self.A * np.exp(x / self.B) + self.k * g(x)) * v.niw \
                     - self.K * g(x) * (agent.actual_speed * tiw) * tiw
        return wall_force

    def move_agents(self):
        for _, agent in self.agents_hashmap.items():
            agent.time += 1
            agent.actual_speed += agent.force * self.time_gap / agent.m
            agent.position += agent.actual_speed * self.time_gap

    def remove_agents(self):
        for name in list(self.agents_hashmap.keys()):
            if self.agents_hashmap[name].position[1] < -0.1:
                del self.agents_hashmap[name]

    def change_agents_destination(self):
        pass
